# -*- coding: utf-8 -*-
# @Time    : 2025/4/6 21:45
# @Author  : sjh
# @Site    : 
# @File    : Stereo.py
# @Comment :
import cv2
import numpy as np
import json
import glob
import os
from stereo.utils.config import Config
# np.set_printoptions(suppress=True)
class stereoRectify:
    def __init__(self, resolution=(640, 480)):
        self.image_size = Config.image_size   # 图像尺寸
        self.initialized = False
        self.initialize(resolution)

    def initialize(self, resolution):
        if self.initialized:
            return
        
        k_left, dist_left, k_right, dist_right, R, T, Q = self.load_calibration_from_json(r'stereo/calib/cali_circle.json')
        self.k_left, self.k_right, self.R, self.T, self.Q, self.dist_left, self.dist_right = np.array(k_left), np.array(k_right), np.array(R), np.array(T), np.array(Q), np.array(dist_left), np.array(dist_right)
        self.map1x, self.map1y, self.map2x, self.map2y, _ = self.stereo_rectify_with_distortion(
            self.k_left, self.dist_left, self.k_right, self.dist_right, 
            self.R, self.T, self.image_size
        )
        self.calculate_fx_fy_cx_cy_baseline(Q)

        ori_width, ori_height = self.image_size  # 原始图像
        scale_x = resolution[0] / ori_width  # 宽度缩放比例
        scale_y = resolution[1] / ori_height  # 高度缩放比例
        self.reset_calib(scale_x, scale_y)  # 调整内参
        print(self.fx, self.fy, self.cx, self.cy, self.baseline)
        self.initialized = True
        
        
    # 计算有畸变立体校正变换
    def stereo_rectify_with_distortion(self, K_left, dist_left, K_right, dist_right, R, T, newimage_size=None):

        # 图像尺寸
        gdc_width, gdc_height = self.image_size    # 原始图像尺寸
        out_width, out_height = Config.image_size     # 校正后图像尺寸

        # 调用stereoRectify
        flags = cv2.CALIB_ZERO_DISPARITY    # 强制校正后主点垂直对齐
        alpha = 0                           # 自动裁剪黑边
        
        R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(
                cameraMatrix1=K_left, # 左相机内参
                distCoeffs1=dist_left, # 左相机畸变系数
                cameraMatrix2=K_right, # 右相机内参
                distCoeffs2=dist_right, # 右相机畸变系数
                imageSize=(gdc_width, gdc_height), # 原始图像尺寸
                R=R, # 旋转矩阵 右到左的旋转矩阵
                T=T, # 平移矩阵 右到左的平移向量
                flags=flags, # 标志
                alpha=alpha, # 自动裁剪黑边
                newImageSize=(out_width, out_height) # 校正后图像尺寸
            )
        # 2. 计算映射表（考虑畸变）
        map1x, map1y = cv2.initUndistortRectifyMap(K_left, dist_left, R1, P1, self.image_size, cv2.CV_32FC1)
        map2x, map2y = cv2.initUndistortRectifyMap(K_right, dist_right, R2, P2, self.image_size, cv2.CV_32FC1)
        
        return map1x, map1y, map2x, map2y, Q
    
    def apply_rectification(self, left_img, right_img, map1x, map1y, map2x, map2y):
        
        # 1. 重映射
        rectified_left = cv2.remap(left_img, map1x, map1y, cv2.INTER_LINEAR)
        rectified_right = cv2.remap(right_img, map2x, map2y, cv2.INTER_LINEAR)
        
        return rectified_left, rectified_right
    # 计算fx,fy,cx,cy,baseline
    def calculate_fx_fy_cx_cy_baseline(self, Q):
        Q = np.array(Q)
        self.cx = -Q[0, 3]
        self.cy = -Q[1, 3]
        self.fx = Q[2, 3]
        self.fy = Q[2, 3]
        self.baseline = abs(1.0 / Q[3, 2]) * 1000

    def reset_calib(self, scale_x, scale_y):
        self.fx, self.fy, self.cx, self.cy = self.fx * scale_x, self.fy * scale_y, self.cx * scale_x, self.cy * scale_y
        self.depth_cam_matrix = np.array([[self.fx, 0, self.cx],
                                          [0, self.fy, self.cy],
                                          [0, 0, 1]])
        print(self.depth_cam_matrix)


    def apply_rectification_single(self, left_img, map1x, map1y, k_left=None, dist_left=None):
        # 1. 畸变校正
        if k_left is not None and dist_left is not None:
            # 如果提供了相机内参和畸变系数，则进行畸变校正
            left_img = cv2.undistort(left_img, k_left, dist_left)
        # 1. 重映射
        rectified_left = cv2.remap(left_img, map1x, map1y, cv2.INTER_LINEAR)
        
        return rectified_left

    # 读取参数
    def load_calibration_from_json(self, file_path):
        with open(file_path, 'r') as json_file:
            calibration_data = json.load(json_file)
        k_left, dist_left, k_right, dist_right, R, T, Q = calibration_data['k_left'], calibration_data['dist_left'], calibration_data['k_right'], calibration_data['dist_right'], calibration_data['R'], calibration_data['T'], calibration_data['Q']
        return k_left, dist_left, k_right, dist_right, R, T, Q
    
    def rectified_image(self, left_image, right_image):
        rectified_left, rectified_right = self.apply_rectification(
            left_image, right_image, self.map1x, self.map1y, self.map2x, self.map2y
        )
        
        # rectified_left = cv2.resize(rectified_left, (640, 480))
        # rectified_right = cv2.resize(rectified_right, (640, 480))
        # rectified_left = rectified_left[:480, :640]
        # rectified_right = rectified_right[:480, :640]
        return rectified_left, rectified_right
    def disparity_to_depth(self, disp, focal_length, baseline):
        """ 视差图转换为深度图 """
        depth = (focal_length * baseline) / (disp + 1e-8)
        return depth


